import gradio as gr
import paddle
import paddle.nn.functional as F
from paddle.vision import transforms
from PIL import Image
import json
import numpy as np

# 确保中文显示正常
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定黑体（SimHei）
plt.rcParams['axes.unicode_minus'] = False  # 解决负号乱码
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

# 垃圾类别标签映射
garbage_classes = {
    0: "有害垃圾",
    1: "可回收物",
    2: "厨余垃圾",
    3: "其他垃圾"
}


# 加载模型配置和参数
def load_model(model_config_path, model_params_path):
    # 创建模型（这里需要根据实际模型结构调整）
    # 定义与训练时相同的模型结构
    class GarbageClassifier(paddle.nn.Layer):
        def __init__(self, num_classes=4):
            super(GarbageClassifier, self).__init__()
            # 这里简化了网络结构，实际应根据模型配置构建
            self.features = paddle.nn.Sequential(
                paddle.nn.Conv2D(3, 64, kernel_size=3, stride=2, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.MaxPool2D(kernel_size=2, stride=2),
                paddle.nn.Conv2D(64, 128, kernel_size=3, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.MaxPool2D(kernel_size=2, stride=2),
                paddle.nn.Conv2D(128, 256, kernel_size=3, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.Conv2D(256, 256, kernel_size=3, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.MaxPool2D(kernel_size=2, stride=2),
            )
            self.avgpool = paddle.nn.AdaptiveAvgPool2D((7, 7))
            self.classifier = paddle.nn.Sequential(
                paddle.nn.Linear(256 * 7 * 7, 512),
                paddle.nn.ReLU(),
                paddle.nn.Dropout(0.5),
                paddle.nn.Linear(512, num_classes),
            )

        def forward(self, x):
            x = self.features(x)
            x = self.avgpool(x)
            x = paddle.flatten(x, 1)
            x = self.classifier(x)
            return x

    model = GarbageClassifier()

    # 加载模型参数
    model_state_dict = paddle.load(model_params_path)
    model.set_state_dict(model_state_dict)
    model.eval()

    return model


def load_model2(model_config_path, model_params_path):
    # 创建模型结构
    class GarbageClassifier(paddle.nn.Layer):
        def __init__(self, num_classes=4):
            super(GarbageClassifier, self).__init__()
            self.features = paddle.nn.Sequential(
                paddle.nn.Conv2D(3, 64, kernel_size=3, stride=2, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.MaxPool2D(kernel_size=2, stride=2),
                paddle.nn.Conv2D(64, 128, kernel_size=3, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.MaxPool2D(kernel_size=2, stride=2),
                paddle.nn.Conv2D(128, 256, kernel_size=3, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.Conv2D(256, 256, kernel_size=3, padding=1),
                paddle.nn.ReLU(),
                paddle.nn.MaxPool2D(kernel_size=2, stride=2),
            )
            self.avgpool = paddle.nn.AdaptiveAvgPool2D((7, 7))
            self.classifier = paddle.nn.Sequential(
                paddle.nn.Linear(256 * 7 * 7, 512),
                paddle.nn.ReLU(),
                paddle.nn.Dropout(0.5),
                paddle.nn.Linear(512, num_classes),
            )

        def forward(self, x):
            x = self.features(x)
            x = self.avgpool(x)
            x = paddle.flatten(x, 1)
            x = self.classifier(x)
            return x

    model = GarbageClassifier()

    # 加载模型参数
    loaded_params = paddle.load(model_params_path)

    # 调试输出
    print(f"加载对象类型: {type(loaded_params)}")

    if isinstance(loaded_params, dict):
        # 如果是字典，正常加载
        model.set_state_dict(loaded_params)
    elif isinstance(loaded_params, paddle.Tensor):
        # 如果是张量，可能需要特殊处理
        print("警告：加载的是张量而非参数字典！")
        print(f"张量形状: {loaded_params.shape}")
        # 这里需要根据具体情况处理，可能需要重新训练或导出正确的模型
    else:
        print(f"不支持的参数类型: {type(loaded_params)}")

    model.eval()
    return model


# 预处理图像
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return transform(image).unsqueeze(0)


# 预测函数
def predict(image):
    if image is None:
        return "请上传一张图片", None

    # 转换为PIL图像
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image.astype('uint8'))

    # 预处理
    input_tensor = preprocess_image(image)

    # 模型预测
    with paddle.no_grad():
        output = model(input_tensor)
        probs = F.softmax(output, axis=1)
        top_prob, top_class = paddle.topk(probs, k=1)

    # 获取预测结果
    predicted_class = int(top_class.numpy()[0][0])
    confidence = float(top_prob.numpy()[0][0]) * 100

    # 返回结果
    result = f"预测类别: {garbage_classes[predicted_class]}\n置信度: {confidence:.2f}%"

    # 生成概率分布图
    fig = plt.figure(figsize=(8, 4))
    plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
    plt.bar(garbage_classes.values(), probs.numpy()[0] * 100, color='skyblue')
    plt.ylim(0, 100)
    plt.ylabel('置信度 (%)')
    plt.title('各类别的预测概率')
    plt.xticks(rotation=45)
    plt.tight_layout()

    return result, fig


# 创建Gradio界面
def create_interface():
    with gr.Blocks(title="垃圾图片分类系统", theme=gr.themes.Soft()) as interface:
        gr.Markdown("### 垃圾图片智能分类系统")
        gr.Markdown("上传一张图片，系统将自动识别该垃圾属于哪一类（有害垃圾、可回收物、厨余垃圾、其他垃圾）")

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="上传图片", type="numpy")
                predict_btn = gr.Button("开始分类", variant="primary")

            with gr.Column():
                output_text = gr.Textbox(label="分类结果")
                output_plot = gr.Plot(label="概率分布")

        examples = [
            # ["examples/battery.jpg"]
        ]

        gr.Examples(
            examples=examples,
            inputs=input_image,
            outputs=[output_text, output_plot],
            fn=predict,
            cache_examples=False,
        )

        predict_btn.click(
            fn=predict,
            inputs=input_image,
            outputs=[output_text, output_plot]
        )

        gr.Markdown("""
        ### 关于垃圾分类
        - **有害垃圾**：含有有害物质，需要特殊安全处理的垃圾，如电池、过期药品等。
        - **可回收物**：适宜回收和资源利用的垃圾，如废纸、塑料、金属等。
        - **厨余垃圾**：易腐烂的生物质废弃物，如剩菜剩饭、果皮等。
        - **其他垃圾**：除上述类别外的其他生活垃圾，如砖瓦陶瓷、尘土等。
        """)

    return interface


# 主函数
if __name__ == "__main__":
    # 加载模型
    # 注意：根据实际情况修改模型路径
    model = load_model2("Model/Garbage.json", "Model/Garbage.pdiparams")

    # 创建界面
    interface = create_interface()

    # 启动应用
    interface.launch(share=True)